/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2024-2024.
 * Description: RTOS Process Context for RPC
 * Create: 2024-03-11
 */
#include <linux/sched.h>
#include <linux/miscdevice.h>
#include <linux/task_struct_extend.h>
#include <linux/list.h>
#include <linux/slab.h>
#include <linux/bug.h>
#include <linux/kref.h>
#include <linux/scs.h>
#include <linux/compiler.h>
#include <linux/capability.h>
#include <linux/uaccess.h>
#include <linux/sched/signal.h>
#include <linux/atomic.h>
#include <linux/rtos_rtpc_rpc.h>

#include <asm/rtos_rtpc.h>

#include "../kernel/sched/sched.h"

#include "rtpc.h"

enum rpc_cmd {
	CMD_BASE = 0xff2d0000,
	CMD_SERVER_INIT,
	CMD_SERVER_EXIT,
	CMD_SERVER_THREAD_BIND,
	CMD_SERVER_THREAD_UNBIND,
	CMD_SERVER_THREAD_DESTROY,
	CMD_CLIENT_INIT,
	CMD_CLIENT_EXIT,
	CMD_CLIENT_THREAD_INIT,
	CMD_CLIENT_THREAD_EXIT,
	CMD_CLIENT_THREAD_BIND,
	CMD_CLIENT_THREAD_UNBIND,
	CMD_DO_CALL,
	CMD_WAIT_AND_RET_CALL,
	CMD_MAX
};

typedef int (*rpc_cmd_func_t)(struct file *filp, unsigned long arg, struct task_struct *tsk);

struct rpc_cmd_node_s {
	unsigned long cmd;
	rpc_cmd_func_t cmd_func;
};

unsigned int sysctl_rpc_kernel_enable;

static LIST_HEAD(service_info_list);
static rtpc_spinlock_t service_info_list_lock;

static void service_info_release(struct kref *ref)
{
	struct rpc_service_info *const service_info =
		container_of(ref, struct rpc_service_info, ref);

	list_del(&service_info->node);

	kfree(service_info);
}

static void service_info_kref_put(struct rpc_service_info *service_info)
{
	unsigned long flags;

	rtpc_spin_lock_irqsave(&service_info_list_lock, flags);
	kref_put(&service_info->ref, service_info_release);
	rtpc_spin_unlock_irqrestore(&service_info_list_lock, flags);
}

static inline void init_session_id(struct rpc_service_info *const service_info,
		const unsigned int session_id)
{
	/*
	 * session_id field should be written only at service_info initialization.
	 * So we set session_id to const in struct rpc_service_info.
	 *
	 * Forcibly changes here to writable for initialization.
	 */
	*(unsigned int *)&service_info->session_id = session_id;
}

static struct rpc_service_info *rpc_find_service_info(const unsigned int session_id)
{
	struct rpc_service_info *service_info;

	list_for_each_entry(service_info, &service_info_list, node) {
		if (service_info->session_id == session_id)
			return service_info;
	}

	return NULL;
}

static struct rpc_service_info *rpc_find_get_service_info(const unsigned int session_id)
{
	struct rpc_service_info *service_info;
	unsigned long flags;

	rtpc_spin_lock_irqsave(&service_info_list_lock, flags);
	service_info = rpc_find_service_info(session_id);
	if (service_info) {
		BUG_ON(kref_read(&service_info->ref) == 0);
		kref_get(&service_info->ref);
	}
	rtpc_spin_unlock_irqrestore(&service_info_list_lock, flags);

	return service_info;
}

static int rpc_add_service(struct rpc_service_info *const service_info)
{
	unsigned long flags;
	int ret = 0;

	rtpc_spin_lock_irqsave(&service_info_list_lock, flags);

	if (rpc_find_service_info(service_info->session_id))
		ret = -1;
	else
		list_add(&service_info->node, &service_info_list);

	rtpc_spin_unlock_irqrestore(&service_info_list_lock, flags);

	return ret;
}

void rtpc_exit(struct task_struct *const tsk)
{
	struct rpc_service_info *const service_info = rpc_get_service_info_self(tsk);
	struct task_struct *calling_client_task;
	struct task_struct *server_task;
	unsigned long flags;

	if (!service_info)
		return;

	spin_lock_irqsave(&tsk->sighand->siglock, flags);
	rtpc_spin_lock(&service_info->lock);

	calling_client_task = service_info->calling_client_task;
	server_task = service_info->server_task;

#ifdef CONFIG_RTOS_RTPC_RPC_DEBUG_BUILD
	if (service_info->need_exit & (unsigned int)EXIT_FIXUP)
		BUG_ON(!service_info->server_need_wake || !service_info->calling_client_need_wake);
#endif

	if (calling_client_task && !service_info->fixup_done) {
		/* Now in a state that requires fixup */
#ifdef CONFIG_RTOS_RTPC_RPC_DEBUG_BUILD
		/* current is not the calling_client or server, is another client */
		BUG_ON(tsk == calling_client_task);
		BUG_ON(tsk == server_task);
		BUG_ON(!server_task);

		/* need_exit normal flag should not be set */
		BUG_ON(service_info->need_exit &
		       (~((unsigned int)KILL_FIXUP | (unsigned int)EXIT_FIXUP)));

		service_info->need_exit |= (unsigned int)EXIT_FIXUP;
#else
		service_info->need_exit = true;
#endif
		rpc_notice("%d/%s exit in fixup\n", tsk->pid, tsk->comm);
		service_info->server_need_wake = true;
		service_info->calling_client_need_wake = true;
	} else {
#ifdef CONFIG_RTOS_RTPC_RPC_DEBUG_BUILD
		if (service_info->need_exit &
		    ((unsigned int)EXIT_NORMAL | (unsigned int)EXIT_FIXUP))
			BUG_ON(calling_client_task);
		BUG_ON(service_info->calling_client_need_wake && calling_client_task);

		service_info->need_exit |= (unsigned int)EXIT_NORMAL;
#else
		service_info->need_exit = true;
#endif

		rpc_notice("%d/%s exit not in fixup\n", tsk->pid, tsk->comm);

		if (tsk == server_task) {
			/*
			 * We must use WRITE_ONCE here,
			 * because arm32 may read with sercive_info no locked
			 */
			WRITE_ONCE(service_info->server_task, NULL);
		} else if (server_task) {
#ifdef CONFIG_RTOS_RTPC_RPC_DEBUG_BUILD
			if (wake_up_state(server_task, TASK_RTPC_UNWAKEABLE))
				BUG_ON(calling_client_task);
#else
			wake_up_state(server_task, TASK_RTPC_UNWAKEABLE);
#endif
		}

		if (calling_client_task) {
#ifdef CONFIG_RTOS_RTPC_RPC_DEBUG_BUILD
			BUG_ON(tsk == calling_client_task);
			BUG_ON(!server_task);
#endif

			WRITE_ONCE(service_info->calling_client_task, NULL);
			wake_up_state(calling_client_task, TASK_RTPC_UNWAKEABLE);
		}
	}

	rtpc_spin_unlock(&service_info->lock);

	/*
	 * we must set current->service_info with siglocked
	 * because other thread sending signal will read p->service_info
	 */
	WRITE_ONCE(task_to_rtos_task(tsk)->rpc_task.service_info, NULL);

	spin_unlock_irqrestore(&tsk->sighand->siglock, flags);
	service_info_kref_put(service_info);
}

static int rpc_server_init(struct file *filp, unsigned long arg, struct task_struct *tsk)
{
	if (!capable(CAP_SYS_ADMIN))
		return -EACCES;

	return 0;
}

static int rpc_server_thread_bind(struct file *filp, unsigned long arg, struct task_struct *tsk)
{
	struct rpc_service_info *service_info;
	unsigned int session_id;

	if (!capable(CAP_SYS_ADMIN))
		return -EACCES;

	if (task_to_rtos_task(tsk)->rpc_task.has_init_as_client) {
		rpc_err("%d/%s already init as client\n", tsk->pid, tsk->comm);
		return -EINVAL;
	}

	if (get_user(session_id, (unsigned int __user *)arg))
		return -EFAULT;

	if (rpc_get_service_info_self(tsk)) {
		rpc_err("%d/%s already bind\n", tsk->pid, tsk->comm);
		return -EEXIST;
	}

	service_info = kcalloc(1, sizeof(struct rpc_service_info), GFP_KERNEL);
	if (!service_info)
		return -ENOMEM;

	init_session_id(service_info, session_id);
	service_info->server_task = tsk;

	if (unlikely(!fair_policy(tsk->policy))) {
		rpc_err("[%d/%s] only support fair sched\n", tsk->pid, tsk->comm);
		kfree(service_info);
		return -EEXIST;
	}

	kref_init(&service_info->ref);
	if (rpc_add_service(service_info)) {
		rpc_err("create already with session_id = %u\n", session_id);
		kfree(service_info);
		return -EEXIST;
	}

	/*
	 * We must use smp_store_release to ensure that service_info->lock is initialized first,
	 * then set task->service_info
	 *
	 * In signal sending, may read task->service_info, and try lock
	 */
	smp_store_release(&task_to_rtos_task(tsk)->rpc_task.service_info, service_info);
	/* Start from here, if server be killed, will set need exit */

	return 0;
}

static int rpc_client_init(struct file *filp, unsigned long arg, struct task_struct *tsk)
{
	if (!capable(CAP_SYS_ADMIN))
		return -EACCES;

	return 0;
}

static int rpc_client_thread_init(struct file *filp, unsigned long arg, struct task_struct *tsk)
{
	if (!capable(CAP_SYS_ADMIN))
		return -EACCES;

	if (task_to_rtos_task(tsk)->rpc_task.has_init_as_client)
		return -EEXIST;

	if (rpc_get_service_info_self(tsk)) {
		rpc_err("%d/%s already bind as server\n", tsk->pid, tsk->comm);
		return -EINVAL;
	}

	task_to_rtos_task(tsk)->rpc_task.has_init_as_client = true;

	return 0;
}

static int rpc_client_thread_bind(struct file *filp, unsigned long arg, struct task_struct *tsk)
{
	int ret = 0;
	struct task_struct *server_task;
	struct rpc_service_info *service_info;
	unsigned int session_id;
	unsigned long flags;

	if (!task_to_rtos_task(tsk)->rpc_task.has_init_as_client)
		return -EINVAL;

	if (get_user(session_id, (unsigned int __user *)arg))
		return -EFAULT;

	if (rpc_get_service_info_self(tsk)) {
		rpc_err("%d/%s already bind\n", tsk->pid, tsk->comm);
		return -EEXIST;
	}

	service_info = rpc_find_get_service_info(session_id);
	if (!service_info) {
		rpc_err("%d/%s not found service_info with session_id = %u\n",
				tsk->pid, tsk->comm, session_id);
		return -ENOENT;
	}

	rtpc_spin_lock_irqsave(&service_info->lock, flags);

	if (service_info->need_exit) {
		rpc_err("this service need exited\n");
		ret = -ESRCH;
		goto out;
	}

	server_task = service_info->server_task;

	if (same_thread_group(tsk, server_task)) {
		rpc_err("client and server in same process\n");
		ret = -EPERM;
		goto out;
	}

	if (service_info->server_init_state != SERVER_RTPC_SLEPT) {
		if (service_info->server_init_state == SERVER_START_SLEEP) {
			BUG_ON(server_task->state != TASK_RTPC_UNWAKEABLE);

			/*
			 * if Client and Server on the same core
			 * spinloop waiting server_task->on_rq may deadlock
			 * so use EAGAIN
			 */
			if (READ_ONCE(server_task->on_rq))
				goto label_eagain;

			/*
			 * This is copy from try_to_wake_up:kernel/sched/core.c
			 * do spinloop wait for server_task->on_cpu change to 0
			 */
			smp_cond_load_acquire(&server_task->on_cpu, !VAL);

#ifdef CONFIG_SHADOW_CALL_STACK
			task_scs_sp(server_task) = task_scs(server_task);
#endif
			rpc_create_context(server_task);
			service_info->server_init_state = SERVER_RTPC_SLEPT;
		} else {
label_eagain:
			rpc_err("%d/%s Server is not completed init\n", tsk->pid, tsk->comm);
			ret = -EAGAIN;
			goto out;
		}
	}

	if (unlikely(!fair_policy(tsk->policy))) {
		rpc_err("[%d/%s] only support fair sched\n", tsk->pid, tsk->comm);
		ret = -EPERM;
		goto out;
	}

	rtpc_spin_unlock_irqrestore(&service_info->lock, flags);

	WRITE_ONCE(task_to_rtos_task(tsk)->rpc_task.service_info, service_info);

	return 0;
out:
	rtpc_spin_unlock_irqrestore(&service_info->lock, flags);
	service_info_kref_put(service_info);
	return ret;
}

static int rpc_empty_error(struct file *filp, unsigned long arg, struct task_struct *tsk)
{
	rpc_err("this interface cannot be use!\n");
	return -EINVAL;
}

static struct rpc_cmd_node_s rpc_cmd_node[] = {
	{CMD_SERVER_INIT, rpc_server_init},
	{CMD_SERVER_EXIT, rpc_empty_error},
	{CMD_SERVER_THREAD_BIND, rpc_server_thread_bind},
	{CMD_SERVER_THREAD_UNBIND, rpc_empty_error},
	{CMD_SERVER_THREAD_DESTROY, rpc_empty_error},
	{CMD_CLIENT_INIT, rpc_client_init},
	{CMD_CLIENT_EXIT, rpc_empty_error},
	{CMD_CLIENT_THREAD_INIT, rpc_client_thread_init},
	{CMD_CLIENT_THREAD_EXIT, rpc_empty_error},
	{CMD_CLIENT_THREAD_BIND, rpc_client_thread_bind},
	{CMD_CLIENT_THREAD_UNBIND, rpc_empty_error},
	/*
	 * In fast path mode, ioctl with CMD_DO_CALL will be intercepted to do_call_fast in:
	 * entry.S: kernel_entry
	 * rtos_rtpc_rpc.c: rpc_do_exception
	 *
	 * It will never run to rpc_get_cmd_node below.
	 */
	{CMD_DO_CALL, rpc_empty_error},
	{CMD_WAIT_AND_RET_CALL, rpc_empty_error},
};

static struct rpc_cmd_node_s *rpc_get_cmd_node(unsigned int cmd)
{
	size_t i;
	const size_t cmd_count = sizeof(rpc_cmd_node) / sizeof(struct rpc_cmd_node_s);

	for (i = 0; i < cmd_count; i++) {
		if (rpc_cmd_node[i].cmd == cmd)
			return &rpc_cmd_node[i];
	}

	return NULL;
}

static long rpc_ioctl(struct file *filp, unsigned int cmd, unsigned long arg)
{
	struct task_struct *const tsk = current;
	struct rpc_cmd_node_s *const cmd_node = rpc_get_cmd_node(cmd);

	if (sysctl_rpc_kernel_enable != RPC_KERNEL_ENABLE_FAST) {
		rpc_err("rpc kernel is not enable!\n");
		return -EPERM;
	}

	if (!cmd_node) {
		rpc_err("task = %s/%d, cmd = 0x%x is invalid\n", tsk->comm, tsk->pid, cmd);
		return -EINVAL;
	}

	return cmd_node->cmd_func(filp, arg, tsk);
}

static const struct file_operations rpc_fops = {
	.owner          = THIS_MODULE,
	.unlocked_ioctl = rpc_ioctl,
	.compat_ioctl   = rpc_ioctl,
};

static struct miscdevice rpc_dev = {
	MISC_DYNAMIC_MINOR,
	"rtos_rpc",
	&rpc_fops
};

static int __init rpc_init(void)
{
	int ret;

	ret = misc_register(&rpc_dev);
	if (ret) {
		rpc_err("register device failed\n");
		return -1;
	}

	pr_info("RTOS Process Context for RPC init success\n");

	return 0;
}

late_initcall(rpc_init);
